package com.yamaikeji.common.utils.security;

/**
 * AES 加密解密工具类
 * @author 桑伟杰
 * @date 2025-04-08
 */
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

public class AesUtils {

    // 默认加密算法
    private static final String ALGORITHM = "AES";
    // 默认加密模式
    private static final String DEFAULT_MODE = "AES/CBC/PKCS5Padding";
    // 默认字符编码
    private static final String CHARSET = "UTF-8";
    // IV 默认长度
    private static final int IV_LENGTH = 16;

    /**
     * AES加密
     * @param content 待加密内容
     * @param key 密钥 (16/24/32字节)
     * @param iv 初始化向量 (16字节)
     * @return Base64编码的加密结果
     */
    public static String encrypt(String content, String key, String iv) {
        return encrypt(content, key, iv, DEFAULT_MODE);
    }

    /**
     * AES加密
     * @param content 待加密内容
     * @param key 密钥 (16/24/32字节)
     * @param iv 初始化向量 (16字节)
     * @param mode 加密模式 (如 "AES/CBC/PKCS5Padding")
     * @return Base64编码的加密结果
     */
    public static String encrypt(String content, String key, String iv, String mode) {
        if (content == null || content.isEmpty()) {
            return "";
        }

        try {
            // 检查密钥长度
            validateKey(key);

            // 创建密钥
            SecretKeySpec secretKey = createSecretKey(key);

            // 创建并初始化密码器
            Cipher cipher = initCipher(Cipher.ENCRYPT_MODE, secretKey, iv, mode);

            // 加密并返回Base64编码结果
            byte[] encrypted = cipher.doFinal(content.getBytes(CHARSET));
            return Base64.getEncoder().encodeToString(encrypted);
        } catch (Exception e) {
            throw new RuntimeException("AES加密失败: " + e.getMessage(), e);
        }
    }

    /**
     * AES解密
     * @param content Base64编码的加密内容
     * @param key 密钥 (16/24/32字节)
     * @param iv 初始化向量 (16字节)
     * @return 解密后的原始字符串
     */
    public static String decrypt(String content, String key, String iv) {
        return decrypt(content, key, iv, DEFAULT_MODE);
    }

    /**
     * AES解密
     * @param content Base64编码的加密内容
     * @param key 密钥 (16/24/32字节)
     * @param iv 初始化向量 (16字节)
     * @param mode 加密模式 (如 "AES/CBC/PKCS5Padding")
     * @return 解密后的原始字符串
     */
    public static String decrypt(String content, String key, String iv, String mode) {
        if (content == null || content.isEmpty()) {
            return "";
        }

        try {
            // 检查密钥长度
            validateKey(key);

            // 创建密钥
            SecretKeySpec secretKey = createSecretKey(key);

            // 创建并初始化密码器
            Cipher cipher = initCipher(Cipher.DECRYPT_MODE, secretKey, iv, mode);

            // 解密并返回结果
            byte[] decoded = Base64.getDecoder().decode(content);
            byte[] decrypted = cipher.doFinal(decoded);
            return new String(decrypted, CHARSET);
        } catch (Exception e) {
            throw new RuntimeException("AES解密失败: " + e.getMessage(), e);
        }
    }

    /**
     * 验证密钥长度
     * @param key 密钥
     */
    private static void validateKey(String key) {
        if (key == null) {
            throw new IllegalArgumentException("密钥不能为空");
        }
        // 如果是Base64编码的密钥，先解码
        if (isBase64(key)) {
            try {
                byte[] keyBytes = Base64.getDecoder().decode(key);
                if (keyBytes.length != 16 && keyBytes.length != 24 && keyBytes.length != 32) {
                    throw new IllegalArgumentException("密钥长度必须是16(AES-128)、24(AES-192)或32(AES-256)字节");
                }
                return;
            } catch (IllegalArgumentException e) {
                // 如果解码失败，继续检查原始字符串
            }
        }

        // 检查原始字符串长度
        int keyLength = key.getBytes(StandardCharsets.UTF_8).length;
        if (keyLength != 16 && keyLength != 24 && keyLength != 32) {
            throw new IllegalArgumentException("密钥长度必须是16(AES-128)、24(AES-192)或32(AES-256)字节");
        }
    }

    /**
     * 判断字符串是否为Base64编码
     * @param str 待检查字符串
     * @return 是否为Base64编码
     */
    private static boolean isBase64(String str) {
        String base64Pattern = "^[A-Za-z0-9+/]*={0,2}$";
        return str.matches(base64Pattern) && str.length() % 4 == 0;
    }

    /**
     * 创建密钥
     * @param key 密钥字符串
     * @return SecretKeySpec对象
     */
    private static SecretKeySpec createSecretKey(String key) {
        byte[] keyBytes;
        if (isBase64(key)) {
            try {
                keyBytes = Base64.getDecoder().decode(key);
            } catch (IllegalArgumentException e) {
                keyBytes = key.getBytes(StandardCharsets.UTF_8);
            }
        } else {
            keyBytes = key.getBytes(StandardCharsets.UTF_8);
        }
        return new SecretKeySpec(keyBytes, ALGORITHM);
    }

    /**
     * 初始化密码器
     * @param mode 加密/解密模式
     * @param secretKey 密钥
     * @param iv 初始化向量
     * @param cipherMode 加密模式
     * @return 初始化后的密码器
     * @throws Exception 初始化异常
     */
    private static Cipher initCipher(int mode, SecretKeySpec secretKey, String iv, String cipherMode) throws Exception {
        Cipher cipher = Cipher.getInstance(cipherMode);
        if (cipherMode.contains("ECB")) {
            cipher.init(mode, secretKey);
        } else {
            if (iv == null || iv.isEmpty()) {
                throw new IllegalArgumentException("非ECB模式下，初始化向量(IV)不能为空");
            }

            byte[] ivBytes;
            if (isBase64(iv)) {
                try {
                    ivBytes = Base64.getDecoder().decode(iv);
                } catch (IllegalArgumentException e) {
                    ivBytes = iv.getBytes(StandardCharsets.UTF_8);
                }
            } else {
                ivBytes = iv.getBytes(StandardCharsets.UTF_8);
            }

            IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
            cipher.init(mode, secretKey, ivParameterSpec);
        }
        return cipher;
    }

    /**
     * 生成随机密钥
     * @param keySize 密钥大小 (128, 192 或 256)
     * @return Base64编码的随机密钥
     */
    public static String generateRandomKey(int keySize) {
        if (keySize != 128 && keySize != 192 && keySize != 256) {
            throw new IllegalArgumentException("密钥大小必须是128、192或256位");
        }

        try {
            byte[] key = new byte[keySize / 8];
            SecureRandom secureRandom = new SecureRandom();
            secureRandom.nextBytes(key);
            return Base64.getEncoder().encodeToString(key);
        } catch (Exception e) {
            throw new RuntimeException("生成随机密钥失败", e);
        }
    }

    /**
     * 生成随机IV
     * @return Base64编码的16字节随机IV
     */
    public static String generateRandomIV() {
        try {
            byte[] iv = new byte[IV_LENGTH];
            SecureRandom secureRandom = new SecureRandom();
            secureRandom.nextBytes(iv);
            return Base64.getEncoder().encodeToString(iv);
        } catch (Exception e) {
            throw new RuntimeException("生成随机IV失败", e);
        }
    }

    public static void main(String[] args) {
        // 生成随机密钥和IV
        String key = generateRandomKey(256);
        String iv = generateRandomIV();
        System.out.println("Key: " + key);
        System.out.println("IV: " + iv);

        // 使用生成的密钥和IV进行加解密测试
        String content = "Hello, AES!";
        String encrypted = encrypt(content, key, iv);
        String decrypted = decrypt(encrypted, key, iv);
        System.out.println("原文: " + content);
        System.out.println("加密: " + encrypted);
        System.out.println("解密: " + decrypted);
    }
}